"""

Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

import pycuda.gpuarray as gpuarray
import numpy as N

from nw_kernels import top_overlay, gpuRho

import kernel_helpers as kh
    
dissim_func = kh.elemwise_inplace('1 - elem') # Func to convert similarity matrix to dissimilarity matrix, in place

def corr(ary, corr_abs=False, threshold=None, dissim=False):
    """Create a correlation matrix from ary. Assumes ary contains sample vectors."""

    M       = kh.aligntobsize(ary, 128)
    M_gpu   = gpuarray.to_gpu(M)
    r, c    = M_gpu.shape

    numsamples, numfeatures = ary.shape

    corrmat = gpuarray.zeros((r, r), dtype=N.float32)

    gpuRho(corrmat, M_gpu, numsamples, numfeatures)

    if corr_abs:
        abs_gpu = kh.elemwise_inplace('fabsf(elem)') # Func to take the abs of an entire matrix, in place
        abs_gpu(corrmat)
    
    if threshold is not None:
        sgn_gpu = kh.elemwise_inplace('(elem >= val) ? 1.0 : 0.0', val=threshold)
        sgn_gpu(corrmat)

    if dissim:
        dissim_func(corrmat)

    c = corrmat.get()[:numsamples][:,:numsamples].astype(N.float32)

    try:
        assert N.isnan(c).any() == False
    except:
        raise ValueError, 'NaNs resulted from correlation computation! Did you normalise your data?'

    return c

def tom(ary, beta=None):
    """Create a TOM matrix from ary. Assumes ary is a similarity matrix."""

    M       = kh.aligntobsize(ary, 128)

    numsamples = ary.shape[0]

    assert ary.shape[0] == ary.shape[1] # Distance matrix?
    
    # Set diagonal to 0
    for i in xrange(M.shape[0]):
        M[i][i] = 0.

    # Transfer distance matrix to GPU
    M_gpu   = gpuarray.to_gpu(M)

    r, c    = M_gpu.shape

    # Convert M_gpu to an adjacency matrix in place

    if beta is not None:
        pow_gpu = kh.elemwise_inplace('pow(elem, val)', val=beta) # Func to take each elem to power val, in place
        pow_gpu(M_gpu) # XXX Add sigmoid

    out = gpuarray.zeros((r,r), dtype=N.float32)

    # Put the topological overlay of M into out
    top_overlay(out, M_gpu)

    # Convert out to a dissimilarity matrix
    dissim_func(out)

    o = out.get()[:numsamples][:,:numsamples].astype(N.float32)

    try:
        assert N.isnan(o).any() == False
    except:
        raise ValueError, 'NaNs resulted from TOM computation!'

    return o
